import os
import cv2
import numpy as np
import argparse
from tqdm import tqdm
import shutil
import scipy.ndimage
from skimage.measure import label
import scipy.ndimage.morphology


def compute_mse_loss(pred, target, trimap):
    error_map = (pred - target) / 255.0
    loss = np.sum((error_map ** 2) * (trimap == 128)) / (np.sum(trimap == 128) + 1e-8)

    # # if test on whole image (Disitinctions-646), please uncomment this line
    # loss = loss = np.sum(error_map ** 2) / (pred.shape[0] * pred.shape[1])

    return loss


def compute_sad_loss(pred, target, trimap):
    error_map = np.abs((pred - target) / 255.0)
    loss = np.sum(error_map * (trimap == 128))

    # # if test on whole image (Disitinctions-646), please uncomment this line
    # loss = np.sum(error_map)

    return loss / 1000, np.sum(trimap == 128) / 1000


def evaluate(args):
    img_names = []
    mse_loss_unknown = []
    sad_loss_unknown = []
    grad_loss_unknown = []
    conn_loss_unknown = []

    bad_case = []

    files = os.listdir(args.label_dir)
    for i, img in tqdm(enumerate(files), total=len(files)):

        if not ((os.path.isfile(os.path.join(args.pred_dir, img)) and
                 os.path.isfile(os.path.join(args.label_dir, img)) and
                 os.path.isfile(os.path.join(args.trimap_dir, img)))):
            print('[{}/{}] "{}" skipping'.format(i, len(os.listdir(args.label_dir)), img))
            continue

        pred = cv2.imread(os.path.join(args.pred_dir, img), 0).astype(np.float32)
        label = cv2.imread(os.path.join(args.label_dir, img), 0).astype(np.float32)
        trimap = cv2.imread(os.path.join(args.trimap_dir, img), 0).astype(np.float32)

        # calculate loss
        mse_loss_unknown_ = compute_mse_loss(pred, label, trimap)
        sad_loss_unknown_ = compute_sad_loss(pred, label, trimap)[0]

        # save for average
        img_names.append(img)

        mse_loss_unknown.append(mse_loss_unknown_)  # mean l2 loss per unknown pixel
        sad_loss_unknown.append(sad_loss_unknown_)  # l1 loss on unknown area

    print('* Unknown Region: MSE:', np.array(mse_loss_unknown).mean(), ' SAD:', np.array(sad_loss_unknown).mean())


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--pred-dir', type=str, required=False, help="output dir")
    parser.add_argument('--label-dir', type=str, default='', help="GT alpha dir")
    parser.add_argument('--trimap-dir', type=str, default='', help="trimap dir")

    dataset_root = os.environ['DATASET_ROOT']
    dataset_adobe = os.path.join(dataset_root, 'adobe/full_dataset/test')
    dataset_d6 = os.path.join(dataset_root, 'd6/Distinctions-646/Test')
    dataset_a5 = os.path.join(dataset_root, 'a5/AIM-500')
    dataset_t4 = os.path.join(dataset_root, 't4/Transparent-460/Test')
    dataset_simd = os.path.join(dataset_root, 'simd/SIMD/test_')


    args = parser.parse_args()
    # args.pred_dir = 'results/c1k/ViTMatte_S_Com'
    # args.pred_dir = 'results/c1k/ViTMatte_B_Com'
    # args.pred_dir = 'results/c1k_hdr/ViTMatte_B_Com' # hdr images
    args.pred_dir = 'results/c1k_half/ViTMatte_B_Com'
    # args.pred_dir = 'results/c1k/ViTMatte_S_DIS'
    # args.pred_dir = 'results/c1k/ViTMatte_B_DIS'
    args.label_dir = '/media/kd/shared/dataset/new_/adobe/full_dataset/test/alpha_copy'
    args.trimap_dir = '/media/kd/shared/dataset/new_/adobe/full_dataset/test/trimaps'

    # args.pred_dir = 'results/dis646/ViTMatte_S_Com'
    # args.pred_dir = 'results/dis646/ViTMatte_B_Com'
    # # args.pred_dir = 'results/dis646/ViTMatte_S_DIS'
    # # args.pred_dir = 'results/dis646/ViTMatte_B_DIS'
    # args.label_dir = '/media/kd/shared/dataset/new_/d6/Distinctions-646/Test/GT'
    # args.trimap_dir = '/media/kd/shared/dataset/new_/d6/Distinctions-646/Test/Trimap'

    # args.pred_dir = 'results/aim500/ViTMatte_S_Com'
    # args.pred_dir = 'results/aim500/ViTMatte_B_Com'
    # args.pred_dir = 'results/aim500_hdr2/ViTMatte_B_Com/predictions/a5' # hdr for official agentmatting paper
    # args.pred_dir = 'results/aim500/ViTMatte_S_DIS'
    # args.label_dir = '/media/kd/shared/dataset/new_/a5/AIM-500/mask'
    # args.trimap_dir = '/media/kd/shared/dataset/new_/a5/AIM-500/trimap'

    # args.pred_dir = 'results/simd/ViTMatte_S_Com'
    # args.pred_dir = 'results/simd_grid22/ViTMatte_S_Com'
    # args.pred_dir = 'results/simd_grid22/ViTMatte_B_Com'
    # args.label_dir = os.path.join(dataset_simd, 'alpha')
    # args.trimap_dir = os.path.join(dataset_simd, 'trimap')

    evaluate(args)
